'''
This file takes the resulting arrf files from GLO_Proj's run and adds AS_Type data to create a test file.
The AS_Type data comes from the classification article writers' file, sn.3m.all.AS.class, and augments
the resulting files from GLO_Proj. All of the measurements arise from GLO's run.
'''

def find_attribute_position(sf,df,attr):
    i=-1
    found=-1
    for line in sf.xreadlines():
        if not df is None: df.write(line)
        if line.startswith('@attribute'):
            i+=1
            if line.startswith('@attribute %s'%attr):
                found=i
        elif line.startswith('@data'):
            break
        
    if -1==found:
        raise "Attribute %s not found in the attribute list"%attr

    return found
    
def create_as_type_dict(types_file):
    types={'1':'EC',
           '2':'STP',
           '3':'LTP',
           '4':'CAHP'}
    f=file(types_file,'r')
    res={}
    for line in f.xreadlines():
        if line.strip()=='': continue
        asn,astype = line.strip().split(' ')
        res[asn]=types[astype]
    return res

def set_as_type(fields,num_pos,type_pos,as_type_dict):
    if not fields[num_pos] in as_type_dict:
        return None
    fields[type_pos]=as_type_dict[fields[num_pos]]
    return ','.join(fields)+'\n'

    
def main():
    directions=['directed','undirected']
    src_file=r"data\%s.AS.all.arff"
    dst_file=r"data\%s.AS.test.arff"
    types_file=r"data\sn.3m.all.AS.class"
    
    for direction in directions:
        as_type_dict=create_as_type_dict(types_file)

        sf=file(src_file%direction,'r')
        df=file(dst_file%direction,'w')
        num_pos=find_attribute_position(sf,None,'asn')
        sf.seek(0)
        type_pos=find_attribute_position(sf,df,'astype')
        
        for line in sf.xreadlines():
            if line.strip()=='':
                df.write(line)
                continue

            fields=line.strip().split(',')
            if fields[type_pos]!='?':
                raise "Field %d is not equal to '?' in line %s"%(type_pos,line)

            line=set_as_type(fields,num_pos,type_pos,as_type_dict)
            if not line is None:
                df.write(line)

        sf.close()
        df.close()
    

if __name__=='__main__':
    main()
